from collections import Counter
from typing import List, Optional

import pyproj

from common.logger import logger
from common.typing_utils import numeric
from common.geometry.geom import Geom
from common.geometry.bbox import BoundingBox, nudge_bbox
from common.geometry.srs import coords_to_epsg_under_utm
from common.geometry.tools import project_geom
from catalog.models.base_stac_item import BaseSTACItem
from catalog.utils import interpret_protocol
import zipfile
import xml.etree.ElementTree as ET
import numpy as np
import numba as nb


def determine_dest_srs(aoi: Optional[Geom], items: List[BaseSTACItem]) -> str:
    """determine destination SRS based on given AOI and a list of items"""
    if aoi is not None:
        lon, lat = aoi.centroid._shapely_geom.x, aoi.centroid._shapely_geom.y
        dest_srs = "EPSG:{}".format(coords_to_epsg_under_utm(lon, lat))
        logger.info(
            f"No `dest_srs` is given, use '{dest_srs}' as the best UTM zone for the given `aoi`."
        )
        return dest_srs

    try:
        counter_epsg_nos = Counter([item.epsg_no for item in items])
        epsg_no = counter_epsg_nos.most_common(1)[0][0]
        dest_srs = "EPSG:{}".format(epsg_no)
        logger.info(
            f"No `dest_srs` is given, use '{dest_srs}' as the most common srs in the involved items."
        )
    except Exception:
        dest_srs = "EPSG:4326"
        logger.info(f"No `dest_srs` is given, use '{dest_srs}' as the best guess.")

    return dest_srs


def determine_dest_res(dest_srs: str, items: List[BaseSTACItem]) -> numeric:
    """Determine destination resolution based on given destination SRS and a list of items"""
    p = pyproj.Proj(dest_srs)
    if p.crs.is_geographic:
        dest_res = items[0].default_resolutions["degree"]
        logger.info(
            f"No `dest_res` is given, use '{dest_res}' as the defined default resolution in degree for that collection."
        )
    else:
        dest_res = items[0].default_resolutions["meter"]
        logger.info(
            f"No `dest_res` is given, use '{dest_res}' as the defined default resolution in meter for that collection."
        )
    return dest_res


def determine_dest_extent(
    extent: Optional[BoundingBox], aoi: Optional[Geom], dest_srs: str, dest_res: numeric
):
    """Determine the destination extent based on given `extent`, `aoi`, destination SRS and resolution."""
    if extent is not None:
        extent_tmp = extent
    elif aoi is not None:
        if dest_srs.upper() != "EPSG:4326":
            projected_aoi = project_geom(aoi, src_epsg="EPSG:4326", dest_epsg=dest_srs)
            extent_tmp = projected_aoi.bounds
        else:
            extent_tmp = aoi.bounds
    else:
        return None

    return nudge_bbox(extent_tmp, dest_res)


def get_s2_l2a_add_boa_offset(item, asset_names):
    """Check if add BOA offset for Sentinel-2 L2A data."""
    # get metadata file path
    # e.g., `/NAS/sentinel2/T49/ +
    # T49SER/S2A_MSIL2A_20220402T030541_N0400_R075_T49SER_20220402T070416.SAFE/ +
    # metadata.zip#MTD_MSIL2A.xml`
    if item.collection_id != "COPERNICUS/SENTINEL2/L2A":
        return None
    fpath_mtd = interpret_protocol(item.info["assets"]["metadata"]["href"])
    fname_meta = "MTD_MSIL2A.xml"
    fpath_meta_zip = fpath_mtd[: -len(f"#{fname_meta}")]
    with zipfile.ZipFile(fpath_meta_zip) as zinfo:
        mtd = zinfo.read(name=fname_meta)
    root = ET.fromstring(mtd)
    child_boa_offset_lst = (
        root[0].find("Product_Image_Characteristics").find("BOA_ADD_OFFSET_VALUES_LIST")
    )
    if child_boa_offset_lst is None:
        return None

    s2_bands = {
        "B01": 0,
        "B02": 1,
        "B03": 2,
        "B04": 3,
        "B05": 4,
        "B06": 5,
        "B07": 6,
        "B08": 7,
        "B8A": 8,
        "B09": 9,
        "B10": 10,
        "B11": 11,
        "B12": 12,
    }
    s2_boa_offsets = {}
    for offset in child_boa_offset_lst.findall("BOA_ADD_OFFSET"):
        s2_boa_offsets[int(offset.attrib["band_id"])] = int(offset.text)

    asset_offsets = []
    for asset in asset_names:
        if asset not in s2_bands:
            asset_offsets.append(0)
        else:
            asset_offsets.append(s2_boa_offsets[s2_bands[asset]])
    return np.array(asset_offsets, dtype=np.int16)


def update_s2_l2a_boa(nda, offsets):
    """Update S2 L2A data by adding the offset."""
    nda = nda.astype(np.int16)
    for i in range(offsets.shape[0]):
        nda[i] = nda[i] + offsets[i]
    nda = np.where(nda >= 0, nda, 0).astype(np.uint16)
    return nda


def is_s2_l2a_collection(items: List[BaseSTACItem]):
    """check if the give items are from the sentinel2-l2a collection"""
    for item in items:
        if item.info["collection"] != "sentinel2-l2a":
            return False
    else:
        return True


@nb.jit(
    nb.uint16[:, :](nb.uint16[:, :, :], nb.uint8[:, :, :]), nopython=True, cache=True
)
def nb_customized_median(nda_values, nda_mask):
    nsrcs, nrows, ncols = nda_values.shape
    nda_res = np.zeros((nrows, ncols), dtype=np.uint16)
    nda_tmp = np.zeros(nsrcs, dtype=np.uint16)
    for i in range(nrows):
        for j in range(ncols):
            available_values = 0
            for k in range(nsrcs):
                if nda_mask[k, i, j] == 0:
                    nda_tmp[available_values] = nda_values[k, i, j]
                    available_values += 1
            if available_values > 0:
                nda_res[i, j] = np.median(nda_tmp[0:available_values])
    return nda_res
